#!/usr/bin/env python3
"""
找出含有D90标注的图片并保存到D90文件夹
"""

import os
import shutil
from pathlib import Path

def find_d90_images():
    """找出含有D90标注的图片"""
    base_path = Path("/root/autodl-tmp/yolo11n_ljy/ultralytics/N-RDD2024/merged_dataset")
    output_dir = Path("/root/autodl-tmp/yolo11n_ljy/ultralytics/D90")
    
    # 创建输出目录
    output_dir.mkdir(exist_ok=True)
    (output_dir / "images").mkdir(exist_ok=True)
    (output_dir / "labels").mkdir(exist_ok=True)
    
    d90_count = 0
    total_checked = 0
    
    print("搜索含有D90标注的图片...")
    
    # 检查训练集和验证集
    for subset in ['train', 'val']:
        labels_dir = base_path / "labels" / subset
        images_dir = base_path / "images" / subset
        
        if not labels_dir.exists():
            print(f"{subset}集标签目录不存在: {labels_dir}")
            continue
            
        print(f"\n检查 {subset} 集...")
        
        # 遍历标签文件
        for label_file in labels_dir.glob("*.txt"):
            total_checked += 1
            
            try:
                with open(label_file, 'r', encoding='utf-8') as f:
                    content = f.read()
                    
                # 检查是否包含类别9 (D90)
                has_d90 = False
                for line in content.strip().split('\n'):
                    if line.strip():
                        parts = line.strip().split()
                        if parts and parts[0] == '9':  # D90对应类别9
                            has_d90 = True
                            break
                
                if has_d90:
                    d90_count += 1
                    
                    # 复制标签文件
                    label_name = f"{subset}_{label_file.name}"
                    shutil.copy2(label_file, output_dir / "labels" / label_name)
                    
                    # 查找对应的图片文件
                    image_name = label_file.stem
                    for ext in ['.jpg', '.jpeg', '.png', '.bmp']:
                        image_file = images_dir / f"{image_name}{ext}"
                        if image_file.exists():
                            image_output_name = f"{subset}_{image_file.name}"
                            shutil.copy2(image_file, output_dir / "images" / image_output_name)
                            print(f"  找到D90图片: {image_output_name}")
                            break
                        
            except Exception as e:
                print(f"  错误读取文件 {label_file}: {e}")
    
    print(f"\n搜索完成!")
    print(f"总共检查了 {total_checked} 个标签文件")
    print(f"找到 {d90_count} 张含有D90标注的图片")
    print(f"文件已保存到: {output_dir}")

if __name__ == "__main__":
    find_d90_images()